# src/dataset.py
import numpy as np
import torch
from torch.utils.data import Dataset
import os


class CriteoDataset(Dataset):
    """
    Criteo数据集加载与预处理 - 懒汉式初始化最终版
    """

    def __init__(self, data_dir, data_type='train', max_vocab_size=100000):
        print(f"初始化 {data_type} 数据集... (实际数据将在worker中加载)")

        self.label_path = os.path.join(data_dir, f'{data_type}_labels.npy')
        self.int_path = os.path.join(data_dir, f'{data_type}_int_features.npy')
        self.cat_path = os.path.join(data_dir, f'{data_type}_cat_features.npy')

        # 直接加载并获取维度信息
        temp_int_data = np.load(self.int_path, mmap_mode='r')
        self.num_samples = temp_int_data.shape[0]
        self.int_dims = temp_int_data.shape[1]

        temp_cat_data = np.load(self.cat_path, mmap_mode='r')
        num_cat_features = temp_cat_data.shape[1]

        self.cat_dims = [max_vocab_size] * num_cat_features

        self.labels = None
        self.int_features = None
        self.cat_features = None

        print(f"初始化完成: {self.num_samples} 条数据")
        print(f"整数特征维度: {self.int_dims}")
        print(f"分类特征数量: {num_cat_features}")

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        if self.labels is None:
            self.labels = np.load(self.label_path, mmap_mode='r')
            self.int_features = np.load(self.int_path, mmap_mode='r')
            self.cat_features = np.load(self.cat_path, mmap_mode='r')

        return {
            'int_features': torch.tensor(self.int_features[idx], dtype=torch.float32),
            'cat_features': torch.tensor(self.cat_features[idx], dtype=torch.long),
            'label': torch.tensor(self.labels[idx], dtype=torch.float32)
        }